# -*- coding: utf-8 -*-
#
# Copyright 2017 zhangxp, st
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import pickle
import numpy as np
import pandas as pd
import datetime as dt

from typing import Dict, List, Optional

from datetime import date,datetime, timedelta
from dateutil.parser import parse

from rqalpha.data.trading_dates_mixin import TradingDatesMixin
from rqalpha.data.base_data_source import BaseDataSource
from rqalpha.utils.i18n import gettext as _
from rqalpha.utils.functools import lru_cache, clear_all_cached_functions
from rqalpha.utils.exception import RQInvalidArgument
from rqalpha.environment import Environment
from rqalpha.model.instrument import Instrument
from rqalpha.const import INSTRUMENT_TYPE, TRADING_CALENDAR_TYPE
from rqalpha.utils.datetime_func import (convert_date_to_int, convert_dt_to_int, convert_int_to_date, convert_int_to_datetime)

from rqalpha.data.base_data_source.storage_interface import (AbstractDateSet,
                                AbstractDayBarStore, AbstractDividendStore,
                                AbstractSimpleFactorStore, AbstractInstrumentStore)
from rqalpha.data.base_data_source.storages import (open_h5, DateSet, DayBarStore, DividendStore, FutureDayBarStore, FutureInfoStore, InstrumentStore, 
                                SimpleFactorStore, ExchangeTradingCalendarStore, YieldCurveStore, ShareTransformationStore)


def to_date(day):
    if isinstance(day, dt.datetime):
        return day.date()
    if isinstance(day, dt.date):
        return day

    if isinstance(day, str):
        return parse(day).date()

    return convert_int_to_date(day).date()

    
class DayBarStoreO(AbstractDayBarStore):
    DEFAULT_DTYPE = np.dtype([
        ('datetime', np.uint64),
        ('open', np.float),
        ('close', np.float),
        ('high', np.float),
        ('low', np.float),
        ('volume', np.float),
    ])

    def __init__(self, path):
        if not os.path.exists(path):
            raise FileExistsError("File {} not exist, please update bundle.".format(path))
        self._path = path
        self._h5 = open_h5(path)

    def __del__(self):
        self._h5.close()

    def get_bars(self, order_book_id):
        try:
            return self._h5[order_book_id][:]
        except KeyError:
            return np.empty(0, dtype=self.DEFAULT_DTYPE)

    def get_date_range(self, order_book_id):
        try:
            data = self._h5[order_book_id]
            return data[0]['datetime'], data[-1]['datetime']
        except KeyError:
            return 20050104, 20050104


class FutureDayBarStoreO(DayBarStoreO):
    DEFAULT_DTYPE = np.dtype(DayBarStoreO.DEFAULT_DTYPE.descr + [("open_interest", '<f8')])


class DividendStoreO(AbstractDividendStore):
    def __init__(self, path):
        self._path = path
        self._h5 = open_h5(path)

    def __del__(self):
        self._h5.close()

    def get_dividend(self, order_book_id):
        try:
            return self._h5[order_book_id][:]
        except KeyError:
            return None


class SimpleFactorStoreO(AbstractSimpleFactorStore):
    def __init__(self, path, is_open=False):
        self._path = path
        self._h5 = open_h5(path)

    def __del__(self):
        self._h5.close()

    def get_factors(self, order_book_id):
        try:
            return self._h5[order_book_id][:]
        except KeyError:
            return None


class DateSetO(AbstractDateSet):
    def __init__(self, f):
        self._f = f
        self._h5 = open_h5(f)

    def __del__(self):
        self._h5.close()

    @lru_cache(None)
    def get_days(self, order_book_id):
        try:
            days = self._h5[order_book_id][:]
        except KeyError:
            return set()
        return set(days.tolist())

    def contains(self, order_book_id, dates):
        date_set = self.get_days(order_book_id)
        if not date_set:
            return None

        def _to_dt_int(d):
            if isinstance(d, (int, np.int64, np.uint64)):
                return int(d // 1000000) if d > 100000000 else int(d)
            else:
                return d.year * 10000 + d.month * 100 + d.day

        return [(_to_dt_int(d) in date_set) for d in dates]

        
class St_Day_KDataSource(BaseDataSource):
    def __init__(self, path, custom_future_info={}, instrument_types=None):
        # 是否保持打开
        if self.__module__ == 'rqalpha_mod_stu_min.data_source':
            if not os.path.exists(path):
                raise RuntimeError('bundle path {} not exist'.format(os.path.abspath(path)))

            def _p(name):
                return os.path.join(path, name)
            # 是否保持打开
            if Environment._env is not None:
                _DateSet = DateSetO
                _DayBarStore = DayBarStoreO
                _DividendStore = DividendStoreO
                _FutureDayBarStore = FutureDayBarStoreO
                _SimpleFactorStore = SimpleFactorStoreO
            else:
                _DateSet = DateSet
                _DayBarStore = DayBarStore
                _DividendStore = DividendStore
                _FutureDayBarStore = FutureDayBarStore
                _SimpleFactorStore = SimpleFactorStore
            #
            funds_day_bar_store = _DayBarStore(_p('funds.h5'))
            self._day_bars = {
                INSTRUMENT_TYPE.CS: _DayBarStore(_p('stocks.h5')),
                INSTRUMENT_TYPE.INDX: _DayBarStore(_p('indexes.h5')),
                INSTRUMENT_TYPE.FUTURE: _FutureDayBarStore(_p('futures.h5')),
                INSTRUMENT_TYPE.ETF: funds_day_bar_store,
                INSTRUMENT_TYPE.LOF: funds_day_bar_store
            }  # type: Dict[INSTRUMENT_TYPE, AbstractDayBarStore]

            self._future_info_store = FutureInfoStore(_p("future_info.json"), custom_future_info)

            self._instruments_stores = {}  # type: Dict[INSTRUMENT_TYPE, AbstractInstrumentStore]
            self._ins_id_or_sym_type_map = {}  # type: Dict[str, INSTRUMENT_TYPE]
            instruments = []
            with open(_p('instruments.pk'), 'rb') as f:
                if instrument_types is not None:
                    for i in pickle.load(f):
                        if i["type"] in instrument_types:
                            if i["type"] == "Future" and Instrument.is_future_continuous_contract(i["order_book_id"]):
                                i["listed_date"] = datetime(1990, 1, 1)
                            instruments.append(Instrument(i, lambda i: self._future_info_store.get_future_info(i)["tick_size"]))
                else:
                    for i in pickle.load(f):
                        if i["type"] == "Future" and Instrument.is_future_continuous_contract(i["order_book_id"]):
                            i["listed_date"] = datetime(1990, 1, 1)
                        instruments.append(Instrument(i, lambda i: self._future_info_store.get_future_info(i)["tick_size"]))
            for ins_type in self.DEFAULT_INS_TYPES:
                self.register_instruments_store(InstrumentStore(instruments, ins_type))


            dividend_store = _DividendStore(_p('dividends.h5'))
            self._dividends = {
                INSTRUMENT_TYPE.CS: dividend_store,
                INSTRUMENT_TYPE.ETF: dividend_store,
                INSTRUMENT_TYPE.LOF: dividend_store,
            }

            self._calendar_providers = {
                TRADING_CALENDAR_TYPE.EXCHANGE: ExchangeTradingCalendarStore(_p("trading_dates.npy"))
            }
            self._tdm = None

            self._yield_curve = YieldCurveStore(_p('yield_curve.h5'))

            split_store = _SimpleFactorStore(_p('split_factor.h5'))
            self._split_factors = {
                INSTRUMENT_TYPE.CS: split_store,
                INSTRUMENT_TYPE.ETF: split_store,
                INSTRUMENT_TYPE.LOF: split_store,
            }
            self._ex_cum_factor = _SimpleFactorStore(_p('ex_cum_factor.h5'))
            self._share_transformation = ShareTransformationStore(_p('share_transformation.json'))

            self._suspend_days = [_DateSet(_p('suspended_days.h5'))]  # type: List[AbstractDateSet]
            self._st_stock_days = _DateSet(_p('st_stock_days.h5'))
        else:
            super(St_Day_KDataSource, self).__init__(path, custom_future_info)

    def __del__(self):
        """"""
        if self._day_bars is not None:
            del self._day_bars
        if self._dividends is not None:
            del self._dividends
        if self._split_factors is not None:
            del self._split_factors
        if self._ex_cum_factor is not None:
            del self._ex_cum_factor
        if self._suspend_days is not None:
            del self._suspend_days
        if self._st_stock_days is not None:
            del self._st_stock_days

    def release(self):
        clear_all_cached_functions()
        #
        del self._day_bars
        self._day_bars = None
        del self._dividends
        self._dividends = None
        del self._split_factors
        self._split_factors = None
        del self._ex_cum_factor
        self._ex_cum_factor = None
        del self._suspend_days
        self._suspend_days = None
        del self._st_stock_days
        self._st_stock_days = None

    def _all_day_bars_of(self, instrument):
        return self._cached_all_day_bars_of(instrument.type, instrument.order_book_id)

    def _filtered_day_bars(self, instrument):
        bars = self._cached_all_day_bars_of(instrument.type, instrument.order_book_id)
        return bars[bars['volume'] > 0]

    @lru_cache(4800)
    def _cached_all_day_bars_of(self, instrument_type, order_book_id):
        return self._day_bars[instrument_type].get_bars(order_book_id)

    def get_dividend(self, instrument):
        return self._get_dividend(instrument.type, instrument.order_book_id)

    @lru_cache(4800)
    def _get_dividend(self, book_type, order_book_id):
        try:
            dividend_store = self._dividends[book_type]
        except KeyError:
            return None

        return dividend_store.get_dividend(order_book_id)

    def get_split(self, instrument):
        return self._get_split(instrument.type, instrument.order_book_id)

    @lru_cache(4800)
    def _get_split(self, book_type, order_book_id):
        try:
            splilt_store = self._split_factors[book_type]
        except KeyError:
            return None

        return splilt_store.get_factors(order_book_id)

    @lru_cache(4800)
    def get_ex_cum_factor(self, order_book_id):
        return self._ex_cum_factor.get_factors(order_book_id)

    def _update_weekly_trading_date_index(self, idx):
        if Environment._env is not None:
            super(St_Day_KDataSource, self)._update_weekly_trading_date_index(idx)
        else:
            if self._tdm is None:
                self._tdm = TradingDatesMixin({t: store.get_trading_calendar() for t, store in self._calendar_providers.items()})
            if self._tdm.is_trading_date(idx):
                return idx
            return self._tdm.get_previous_trading_date(idx)
    
    def resample_week_bars(self, bars, bar_count, fields):
        fields = bars.dtype.names if fields is None else fields
        return super(St_Day_KDataSource, self).resample_week_bars(bars, bar_count, fields)

    def _get_data_range(self, instrument_type, universe):
        #
        universe = [universe] if isinstance(universe, str) else universe
        min_s, max_e = 20991231000000, 20000101000000
        for u in universe:
            s, e = self._day_bars[instrument_type].get_date_range(u)
            min_s = s if s < min_s else min_s
            max_e = e if e > max_e else max_e
        return min_s, max_e

    def available_data_range(self, frequency):
        # FIXME
        from rqalpha.const import DEFAULT_ACCOUNT_TYPE
        accounts = Environment.get_instance().config.base.accounts
        if not (DEFAULT_ACCOUNT_TYPE.STOCK in accounts or DEFAULT_ACCOUNT_TYPE.FUTURE in accounts):
            return date.min, date.max
        if frequency in ['tick', '1d']:
            if 'params' in Environment.get_instance().config and \
                'universe' in Environment.get_instance().config.params:
                if DEFAULT_ACCOUNT_TYPE.STOCK in accounts and DEFAULT_ACCOUNT_TYPE.FUTURE not in accounts:
                    s, e = self._get_data_range(INSTRUMENT_TYPE.CS, Environment.get_instance().config.params.universe)
                    return convert_int_to_date(s).date(), convert_int_to_date(e).date()
                if DEFAULT_ACCOUNT_TYPE.FUTURE in accounts and DEFAULT_ACCOUNT_TYPE.STOCK not in accounts:
                    s, e = self._get_data_range(INSTRUMENT_TYPE.FUTURE, Environment.get_instance().config.params.universe)
                    return convert_int_to_date(s).date(), convert_int_to_date(e).date()
                # all
                s, e = self._day_bars[INSTRUMENT_TYPE.INDX].get_date_range('000001.XSHG')
                return convert_int_to_date(s).date(), convert_int_to_date(e).date()
            elif 'sys_analyser' in Environment.get_instance().config.mod and \
                'benchmark' in Environment.get_instance().config.mod.sys_analyser:
                if not isinstance(Environment.get_instance().config.mod.sys_analyser.benchmark, str):
                    benchmarks = Environment.get_instance().config.mod.sys_analyser.benchmark
                    # 字典
                    for order_book_id, weight in benchmarks.items():
                        benchmark = order_book_id
                        break
                else:
                    benchmark = Environment.get_instance().config.mod.sys_analyser.benchmark
                inst = Environment.get_instance().data_proxy.instrument(benchmark)
                s, e = self._day_bars[inst.type].get_date_range(inst.order_book_id)
                return convert_int_to_date(s).date() + timedelta(days=1), convert_int_to_date(e).date()
            else:
                s, e = self._day_bars[INSTRUMENT_TYPE.INDX].get_date_range('000001.XSHG')
                return convert_int_to_date(s).date(), convert_int_to_date(e).date()


class St_Day_Min_KDataSource(St_Day_KDataSource):
    def __init__(self, path, custom_future_info={}, instrument_types=None):
        # 比父类先初始化，以便父类构造函数里用
        super(St_Day_Min_KDataSource, self).__init__(path, custom_future_info, instrument_types)
        #
        min_path = path.replace('bundle', 'bundle_min') if path.endswith('bundle') else path
        def _p(name):
            return os.path.join(min_path, name)
        # 是否保持打开
        if self.__module__ == 'rqalpha_mod_stu_min.data_source' and Environment._env is not None:
            _FutureDayBarStore = FutureDayBarStoreO
        else:
            _FutureDayBarStore = FutureDayBarStore
        self._cur_freq = '1m'
        if os.path.exists(_p('futures_1m.h5')):
            self._min_bars = self._1min_bars = {
                INSTRUMENT_TYPE.FUTURE: _FutureDayBarStore(_p('futures_1m.h5'))
            }
        else:
            self._min_bars = self._1min_bars = {}
        # 存在才打开，避免异常
        if os.path.exists(_p('futures_5m.h5')):
            self._5min_bars = {
                INSTRUMENT_TYPE.FUTURE: _FutureDayBarStore(_p('futures_5m.h5'))
            }
        else:
            self._5min_bars = {}
        if os.path.exists(_p('futures_10m.h5')):
            self._10min_bars = {
                INSTRUMENT_TYPE.FUTURE: _FutureDayBarStore(_p('futures_10m.h5'))
            }
        else:
            self._10min_bars = {}
        if os.path.exists(_p('futures_15m.h5')):
            self._15min_bars = {
                INSTRUMENT_TYPE.FUTURE: _FutureDayBarStore(_p('futures_15m.h5'))
            }
        else:
            self._15min_bars = {}
        if os.path.exists(_p('futures_30m.h5')):
            self._30min_bars = {
                INSTRUMENT_TYPE.FUTURE: _FutureDayBarStore(_p('futures_30m.h5'))
            }
        else:
            self._30min_bars = {}
        if os.path.exists(_p('futures_60m.h5')):
            self._60min_bars = {
                INSTRUMENT_TYPE.FUTURE: _FutureDayBarStore(_p('futures_60m.h5'))
            }
        else:
            self._60min_bars = {}
        # 缓存加速
        self._current_bars = {}

    def __del__(self):
        """"""
        if self._min_bars is not None:
            del self._min_bars
        if self._1min_bars is not None:
            del self._1min_bars
        if self._5min_bars is not None:
            del self._5min_bars
        if self._10min_bars is not None:
            del self._10min_bars
        if self._15min_bars is not None:
            del self._15min_bars
        if self._30min_bars is not None:
            del self._30min_bars
        if self._60min_bars is not None:
            del self._60min_bars

    def release(self):
        super(St_Day_Min_KDataSource, self).release()
        # 
        del self._min_bars
        self._min_bars = None
        del self._1min_bars
        self._1min_bars = None
        del self._5min_bars
        self._5min_bars = None
        del self._10min_bars
        self._10min_bars = None
        del self._15min_bars
        self._15min_bars = None
        del self._30min_bars
        self._30min_bars = None
        del self._60min_bars
        self._60min_bars = None

    def _set_frequency(self, frequency):
        if frequency == '1m':
            self._cur_freq = '1m'
            self._min_bars = self._1min_bars
        elif frequency == '5m':
            self._cur_freq = '5m'
            self._min_bars = self._5min_bars
        elif frequency == '10m':
            self._cur_freq = '10m'
            self._min_bars = self._10min_bars
        elif frequency == '15m':
            self._cur_freq = '15m'
            self._min_bars = self._15min_bars
        elif frequency == '30m':
            self._cur_freq = '30m'
            self._min_bars = self._30min_bars
        elif frequency == '60m':
            self._cur_freq = '60m'
            self._min_bars = self._60min_bars
        else:
            raise RQInvalidArgument("can not get minute history in %s frequency" % frequency)
        pass

    def available_data_range(self, frequency):
        # FIXME
        from rqalpha.const import DEFAULT_ACCOUNT_TYPE
        accounts = Environment.get_instance().config.base.accounts
        if not (DEFAULT_ACCOUNT_TYPE.STOCK in accounts or DEFAULT_ACCOUNT_TYPE.FUTURE in accounts):
            return date.min, date.max
        if frequency in ['tick', '1d']:
            return super(St_Day_Min_KDataSource, self).available_data_range(frequency)
        elif frequency in ['1m', '5m', '10m', '15m', '30m', '60m']:
            if DEFAULT_ACCOUNT_TYPE.FUTURE in accounts:
                self._set_frequency(frequency)
                if 'params' not in Environment.get_instance().config or \
                    'universe' not in Environment.get_instance().config.params:
                    return date.min, date.max
                s, e = self._get_data_range(INSTRUMENT_TYPE.FUTURE, Environment.get_instance().config.params.universe)
                return convert_int_to_date(s).date(), convert_int_to_date(e).date()

    def instrument(self, sym_or_id):
        return next(iter(self.get_instruments(id_or_syms=[sym_or_id])), None)

    def assure_instrument(self, id_or_ins):
        if isinstance(id_or_ins, Instrument):
            return id_or_ins
        elif isinstance(id_or_ins, str):
            return self.instrument(id_or_ins)
        else:
            raise _(u"unsupported order_book_id type")

    def assure_order_book_id(self, id_or_ins):
        return self.assure_instrument(id_or_ins).order_book_id

    def get_trading_minutes_for(self, instrument, trading_dt):
        if INSTRUMENT_TYPE.FUTURE == instrument.type:
            end_time = trading_dt.replace(hour=15, minute=30, second=0)
            prev_dt = Environment.get_instance().data_proxy.get_previous_trading_date(trading_dt)
            if prev_dt.date() < trading_dt.date(): # 上一个交易日存在且有效
                begin_time = prev_dt.replace(hour=15, minute=30, second=0)
            else:
                begin_time = end_time - timedelta(hours=24)

            bars = self._all_min_bars_of(instrument)
            if len(bars) <= 0:
                return
            dt = np.uint64(convert_dt_to_int(begin_time))
            i_left = bars['datetime'].searchsorted(dt)
            if i_left >= len(bars) or bars['datetime'][i_left] < dt:
                return None
            dt = np.uint64(convert_dt_to_int(end_time))
            i_right = bars['datetime'].searchsorted(dt)
            if i_right >= len(bars) or bars['datetime'][i_right] < dt:
                return None
            # 缓存加速
            self._current_bars[instrument.order_book_id] = {'bars':bars, 'left':i_left, 'right':i_right}
            return bars['datetime'][i_left:i_right]
        else:
            raise NotImplementedError

    def _all_min_bars_of(self, instrument):
        return self._cached_all_min_bars_of(instrument.type, instrument.order_book_id, self._cur_freq)


    @lru_cache(2048)
    def _cached_all_min_bars_of(self, instrument_type, order_book_id, cur_freq):
        assert(cur_freq == self._cur_freq)
        return self._min_bars[instrument_type].get_bars(order_book_id)


    def get_bar(self, instrument, dt, frequency):
        # 
        if frequency == '1d':
            return super(St_Day_Min_KDataSource, self).get_bar(instrument, dt, frequency)
        elif frequency != self._cur_freq:
            raise RQInvalidArgument("can not get minute history in diff frequency")
        # 缓存加速
        bars_dict = self._current_bars[instrument.order_book_id]
        bars = bars_dict['bars'][bars_dict['left']:bars_dict['right']]
        if len(bars) <= 0:
            return
        dt = np.uint64(convert_dt_to_int(dt))
        pos = bars['datetime'].searchsorted(dt)
        if pos >= len(bars) or bars['datetime'][pos] != dt:
            return None

        return bars[pos]


    def history_bars(self, instrument, bar_count, frequency, fields, dt, skip_suspended=True, include_now=False,
                     adjust_type='pre', adjust_orig=None):
        if frequency in ['1d', '1w']:
            return super(St_Day_Min_KDataSource, self).history_bars(instrument, bar_count, frequency, fields, dt, skip_suspended, include_now, adjust_type, adjust_orig)
        elif frequency != self._cur_freq:
            self._set_frequency(frequency)

        if INSTRUMENT_TYPE.FUTURE == instrument.type:
            # 缓存加速
            if instrument.order_book_id in self._current_bars:
                bars_dict = self._current_bars[instrument.order_book_id]
                bars = bars_dict['bars'][bars_dict['left']-bar_count:bars_dict['right']]
            else:
                bars = self._all_min_bars_of(instrument)

            if not self._are_fields_valid(fields, bars.dtype.names):
                raise RQInvalidArgument("invalid fields: {}".format(fields))

            if len(bars) <= 0:
                return bars

            dt = np.uint64(convert_dt_to_int(dt))
            i = bars['datetime'].searchsorted(dt, side='right')
            if bar_count > 0:
                left = i - bar_count if i >= bar_count else 0
            else:
                dt = np.uint64(convert_dt_to_int(adjust_orig))
                left = bars['datetime'].searchsorted(dt, side='left')
            bars = bars[left:i]
            if adjust_type == 'none' or instrument.type in {'Future', 'INDX'}:
                # 期货及指数无需复权
                return bars if fields is None else bars[fields]
            else:
                raise NotImplementedError
        else:
            raise NotImplementedError


if __name__ == '__main__':
    pass